#!/usr/bin/python
# -*- coding: utf-8 -*-

# Copyright (c) 2011
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# Author: Jesus Carrero <j.o.carrero@gmail.com>
# Mountain View, CA

""" Calibrate volatility surface by means of Euler-Lagrange approach.
    Please read the paper:

"""

from numpy import c_, tile, meshgrid, zeros, ones, dot, asarray
from scipy import linalg
from Poisson2DFiniteDifference import Poisson2DFiniteDifference
from scipy.sparse.linalg import spsolve as spsolve
from itertools import izip, count

import logging
import pdb


class CalibVolSurfCCE:#(Poisson2DFiniteDifference):

    """
    Algorithm:
    1- Introduce a function. This will be the initial approximation to the true volatility.
    2- At each point determine '0 ?m?S; n?t?:
    3- Use the Black-Scholes formula to approximate the load.
    4- Use the values obtained in 3 to approximate.
    5- Solve the Poisson equation
    6- Take the function obtained in step 5 and call this '1 : Repeat the above procedure
       from step 3, using '1 in place of '0 to obtain the next approximation. Continue
       until the di erence between successive iterates 'k and 'k+1 is smaller than some
       desired tolerance.
    """

    __slots__ = ['m_price_grid', 'm_time_grid', 'm_strikes', 'm_opt_prices',
                 'm_int_rate', 'm_sigma', 'm_price_eng', 'm_tol', 'm_max_niter',
                 'm_poiss_solver']

    def __init__(self, one2dim=2):
        """ instantiate an object of type volatility calibration."""

        (self.m_price_grid, self.m_time_grid) = (None, None)
        (self.m_strikes, self.m_opt_prices, self.m_int_rate) = (None, None,
                None)
        (self.m_time2_mat, self.m_relax) = (None, 1.0)
        (self.m_sigma, self.m_price_eng, self.m_vega) = (None, None,
                None)

        (self.m_tol, self.m_max_niter) = (1.0e-03, 100)
        self.m_log_progress = logging.getLogger('CalibrateCCE')
        self.m_init_guesss = None
        self.m_load, self.m_time_posi, self.m_price_posi = None, None, None

        self.m_one2dim = one2dim

        if 2 == one2dim:
            self.m_poiss_solver = Poisson2DFiniteDifference()
        else:
            self.m_poiss_solver = None

    def set_simu_domain(self, leftx, botty, rightx, topy):
        """ dirichlet boundary terms. """

        assert 2 == self.m_one2dim
        self.m_poiss_solver.set_simu_domain(leftx, botty, rightx, topy)

    def set_bott_top(self, bott, top):
        """ top bottom bc. """

        assert 2 == self.m_one2dim
        self.m_poiss_solver.set_bott_top(bott, top)

    def set_xlxr(self, left, right):
        """ left and right bc """

        self.m_poiss_solver.set_xlxr(left, right)

    def set_dx_dy(self, delta_x, delta_y=0.):
        """ granularity of the grid in x and y direction. """

        if 2 == self.m_one2dim:
            self.m_poiss_solver.set_dx_dy(delta_x, delta_y)

    def set_calibration_data(self, strikes, time2_mat, prices):
        """
            Calibation data is a triple of Strike Prices,
            Time to Expiration and Observed Derivative prices.
        """

        (self.m_strikes, self.m_time2_mat, self.m_opt_prices) = \
            (asarray(strikes), asarray(time2_mat), asarray(prices))
        self.m_strikes = self.m_strikes.reshape((self.m_strikes.size,
                1))
        self.m_time2_mat = \
            self.m_time2_mat.reshape((self.m_time2_mat.size, 1))
        self.m_opt_prices = \
            self.m_opt_prices.reshape((self.m_opt_prices.size, 1))

    def set_interest_rate(self, interest_rate):
        """ interest rate is assume constant in calibration. """

        self.m_int_rate = float(interest_rate)

    def set_analytical_enginee(self, price_eng, vega=None):
        """ set pricing engine, call, put, digital. Only one
            allow in here. can be easyly extend to handle any engine.
        """

        (self.m_price_eng, self.m_vega) = (price_eng, vega)

    def set_initial_guess(self, init_guess):
        """ assume volatility is constant to start with. """

        self.m_init_guesss = init_guess

    def set_tol_and_niter(self, tol=1.0e-03, max_num_iter=100):
        """ stopping criterias. """

        (self.m_tol, self.m_max_niter) = (tol, max_num_iter)

    def set_relax_param(self, _lambda):
        """ comming from Euler-Lagrange equations. """

        self.m_relax = float(_lambda)

    def __set_up_calibration(self):
        """ prepare everthing for  calibration. """

        assert self.m_strikes.size == self.m_time2_mat.size
        assert self.m_strikes.size == self.m_opt_prices.size
        assert self.m_price_eng is not None
        assert self.m_int_rate is not None

        # Initilize base engine

        solver = self.m_poiss_solver
        solver.set_up_engine()

        # Computational grid

        if 2 == self.m_one2dim:
            (self.m_price_posi, self.m_time_posi) = solver.get_indx_grid()
            (self.m_price_grid, self.m_time_grid) = solver.get_mesh()
        else:
            assert False

        # Genereate initial guess

        if self.m_sigma is None:
            (price_grid, time_grid) = (self.m_price_grid,
                    self.m_time_grid)

            assert self.m_init_guesss is not None
            self.m_sigma = self.m_init_guesss * ones((price_grid.size,
                    time_grid.size), 'float')
            self.m_load = zeros(self.m_sigma.shape)
            self.__update_discrete_load()

        solver.set_discrete_load(self.m_load)

    def __price_contracts(self, posi_in_grid):
        """ price derivatives base on current conditions. """

        prices = zeros(self.m_strikes.shape)
        for (i, k, t2mat) in izip(count(), self.m_strikes,
                              self.m_time2_mat):
            (stock_price, k, t2mat) = ([self.m_price_grid[posi_in_grid[0]]],
                    float(k), float(t2mat) - self.m_time_grid[posi_in_grid[1]])

            if t2mat <= 0:
                prices[i] = 0.
                continue

            prices[i] = self.m_price_eng(stock_price, k, self.m_int_rate,
                    self.m_sigma[posi_in_grid[0], posi_in_grid[1]], t2mat)
        return prices

    def __vegas(self, posi_in_grid):
        """
            Need to define a trust region for numerical derivatieves.
            At this points is just at hoc
        """

        vegas = zeros(self.m_strikes.shape)
        for (i, k, expiry) in izip(count(), self.m_strikes,
                self.m_time2_mat):
            (price_posi, strike, reduced_time) = \
                (self.m_price_grid[posi_in_grid[0]], float(k), float(expiry)
                 - self.m_time_grid[posi_in_grid[1]])
            if reduced_time <= 0:
                vegas[i] = 0.
                continue

            vegas[i] = self.m_vega([price_posi], strike, self.m_int_rate,
                                   self.m_sigma[posi_in_grid[0],
                                   posi_in_grid[1]], reduced_time)
        return vegas

    def __update_discrete_load(self):
        """ assemble load as in the paper. """

        (xgrid, ygrid) = meshgrid(self.m_price_posi[1:-1],
                                  self.m_time_posi[1:-1])

        if 1 == self.m_one2dim:
            (xgrid, ygrid) = (xgrid[0], ygrid[0])

        load = zeros( (xgrid.size, 1), 'float64')
        for (i, posi_in_grid) in izip(count(), c_[xgrid.flatten(),
                             ygrid.flatten()]):
            curr_val_minus_price = self.__price_contracts(posi_in_grid)
            vegas = self.__vegas(posi_in_grid)
            load[i] = self.m_relax * dot(vegas.T, curr_val_minus_price
                    - self.m_opt_prices)

        (nypnts, nxpnts) = xgrid.shape

        self.m_load[1:-1, 1:-1] = -tile(load[0: nxpnts], (nypnts, 1)).reshape((nypnts, nxpnts))
        #self.m_load[1:-1, 1:-1] = load.reshape((nypnts, nxpnts))

    def calibrate(self):
        """ calibrate model. """

        self.__set_up_calibration()

        # Calibration loop

        solver = self.m_poiss_solver
        (M, rhs) = solver.discretize()

        (nxpts, nypts) = (self.m_price_grid.size, self.m_time_grid.size)

        (vbott, vtop, vleft, vright) = solver.get_discretized_bc()

        """
        current_load = solver.post_process(nxpts, nypts,
                                   vbott, vtop, vleft, vright,
                                   self.m_load[1: -1, 1 :-1])

        return (self.m_price_grid, self.m_time_grid, current_load)
        """

        (nxp, nyp) = (int(nxpts - 2), int(nypts - 2))

        (error, iter_num) = (1000.0, 0)
        while iter_num < self.m_max_niter:
            updated_sigma = spsolve(M, rhs)

            # compute error

            error_new = linalg.norm(updated_sigma - self.m_sigma[1:-1,
                                    1:-1].flatten())
            if abs(error_new - error) < self.m_tol:
                break
            else:
                error = error_new
            print " current error. ", error

            # update sigma

            self.m_sigma[1:-1, 1:-1] = updated_sigma.reshape(nyp,
                    nxp).T.copy()

            # compute the new load vector.

            self.__update_discrete_load()
            solver.set_discrete_load(self.m_load)
            """
            current_load = solver.post_process(nxpts, nypts,
                                   vbott, vtop, vleft, vright,
                                   self.m_load[1: -1, 1 :-1])

            return (self.m_price_grid, self.m_time_grid, current_load)
            """
            # recompute rhs.

            rhs = solver.get_rhs()
            iter_num = iter_num + 1

        solu = solver.post_process(nxpts, nypts,
                                 vbott, vtop, vleft, vright,
                                 updated_sigma)

        return (self.m_price_grid, self.m_time_grid, solu)


